import json
import os

import util.common
from util.llm import http_request

MODEL_NAME = os.getenv('MODEL_NAME', "deepseek-chat")
SYSTEM_PROMPT = """你是一名搜索专家，善于使用各种搜索引擎。
你可以使用百度、谷歌等搜索引擎。
我将给出一个问题，你需要理解这个问题，并从问题中提取关键词用于搜索。
提取的关键词可以是多个，关键词使用json array表示。
输出示例：
如果需要输出的内容为"A条文"、"B条文"、"C条文"，则返回的内容为：["A条文","B条文","C条文"]"""

PROMPT = """以下我将给出一个问题，请帮我提取关键词用于搜索引擎搜索。结果用json array表示
问题：{}"""


async def query_rewrite(query, context):
    result = await http_request(
        system_prompt=SYSTEM_PROMPT,
        history=context,
        prompt=PROMPT.format(query),
        # model="hunyuan-turbo"
        model=MODEL_NAME
    )
    try:
        result = util.common.remove_markdown_code_block(result)
        result = json.loads(result)
        result.append(query)
        result = list(set(result))
    except Exception as e:
        print("query_rewrite error:{}".format(e))
        result = [query, ]
    print("query rewrite:{}".format(result))
    return result
